package com.banmoon.datastructure.RedBlackTree;

import lombok.Data;

import java.util.Objects;
import java.util.Optional;

import static java.util.Objects.nonNull;

/**
 * 红黑树
 *
 * @param <K> 键
 * @param <V> 值
 */
public class BRTree<K extends Comparable<K>, V> {
    /**
     * 红黑树颜色-红
     */
    public static final Boolean RED = true;
    /**
     * 红黑树颜色-黑
     */
    public static final Boolean BLACK = false;

    /**
     * 根节点
     */
    public BRTreeNode<K, V> root;

    public  BRTree(K key, V value) {
        root = new BRTreeNode<>(key, value, BLACK);
    }

    /**
     * 中序遍历
     */
    public String middleShow() {
        if (nonNull(root)) {
            return root.middleShow();
        }
        return null;
    }

    /**
     * 左旋 <br />
     * 1、将右子树的父节点 -> 当前节点的父节点 <br />
     * 2、将当前节点的父节点 -> 右子树的左节点 | 右儿子变爸爸，爸爸变左儿子 <br />
     * 3、原先右节点的左子树 -> 改为当前节点的右节点 <br />
     */
    private void leftRotate(BRTreeNode<K, V> node) {
        BRTreeNode<K, V> right = node.right;
        BRTreeNode<K, V> parent = node.parent;
        BRTreeNode<K, V> leftByRight = right.left;
        // 1、将右子树的父节点 -> 当前节点的父节点
        if (nonNull(parent)) {
            right.parent = parent;
            parent.right = right;
        } else {
            right.parent = null;
            this.root = right;
        }
        // 2、将当前节点的父节点 -> 左子树的左节点 | 左儿子变爸爸，爸爸变右儿子
        right.left = node;
        node.parent = right;
        // 3、原先右节点的左子树 -> 改为当前节点的右节点
        node.right = leftByRight;
        if (nonNull(leftByRight)) {
            leftByRight.parent = node;
        }
    }

    /**
     * 右旋 <br />
     * 1、将左子树的父节点 -> 当前节点的父节点 <br />
     * 2、将当前节点的父节点 -> 左子树的右节点 | 左儿子变爸爸，爸爸变右儿子 <br />
     * 3、原先左节点的右子树 -> 改为当前节点的左节点 <br />
     *
     * @param node 当前红黑树
     */
    private void rightRotate(BRTreeNode<K, V> node) {
        BRTreeNode<K, V> left = node.left;
        BRTreeNode<K, V> parent = node.parent;
        BRTreeNode<K, V> rightByLeft = left.right;
        // 1、将左子树的父节点 -> 当前节点的父节点
        if (nonNull(parent)) {
            left.parent = parent;
            parent.left = left;
        } else {
            left.parent = null;
            this.root = left;
        }
        // 2、将当前节点的父节点 -> 左子树的右节点 | 左儿子变爸爸，爸爸变右儿子
        left.right = node;
        node.parent = left;
        // 3、原先左子树的右子树 -> 改为当前节点的左节点
        node.left = rightByLeft;
        if (nonNull(rightByLeft)) {
            rightByLeft.parent = node;
        }
    }

    /**
     * 插入
     * @param key 键
     * @param value 值
     * @return 旧值
     */
    public V add(K key, V value) {
        if (Objects.isNull(this.root)) {
            this.root = new BRTreeNode<>(key, value);
            this.root.toggleColor();
            return null;
        } else {
            return Optional.ofNullable(root.insert(new BRTreeNode<>(key, value), this))
                    .map(BRTreeNode::getValue)
                    .orElse(null);
        }
    }

    @Data
    public static class BRTreeNode<K extends Comparable<K>, V> {

        /**
         * 键
         */
        private K key;
        /**
         * 值
         */
        private V value;

        /**
         * 颜色
         */
        private Boolean color;
        /**
         * 父节点
         */
        private BRTreeNode<K, V> parent;
        /**
         * 左子节点
         */
        private BRTreeNode<K, V> left;
        /**
         * 右子节点
         */
        private BRTreeNode<K, V> right;

        public BRTreeNode(K key, V value) {
            this(key, value, RED);
        }

        public BRTreeNode(K key, V value, Boolean color) {
            this.key = key;
            this.value = value;
            this.color = color;
        }

        /**
         * 中序遍历
         */
        public String middleShow() {
            StringBuilder sb = new StringBuilder();
            if (nonNull(left))
                sb.append(left.middleShow());
//            System.out.println("key：" + key + "，value：" + value);
            System.out.println(String.format("value: %s, 颜色: %s", value, color ? "红" : "黑"));
            sb.append(value + " ");
            if (nonNull(right))
                sb.append(right.middleShow());
            return sb.toString();
        }

        /**
         * 插入
         * @param node 插入的节点
         * @return 旧节点
         */
        public BRTreeNode<K, V> insert(BRTreeNode<K, V> node, BRTree<K, V> tree) {
            K thisKey = this.key;
            K insertKey = node.getKey();
            int compare = thisKey.compareTo(insertKey);
            // 当key值相等，则说明要进行替换
            if (compare == 0) {
                node.parent = this.parent;
                if (Objects.isNull(this.parent)) {
                    tree.root = node;
                }
                if (nonNull(this.left)) {
                    this.left.parent = node;
                    node.setLeft(this.left);
                }
                if (nonNull(this.right)) {
                    this.right.parent = node;
                    node.setRight(this.right);
                }
                // 颜色需要变得和当前节点一样
                node.setColor(this.color);
                return this;
            }
            // 当前key比较大，需要放置左边
            else if (compare > 0) {
                if (nonNull(this.left)) {
                    this.left.insert(node, tree);
                } else {
                    this.left = node;
                    node.parent = this;
                    node.balanceTree(true, tree);
                }
            }
            // 当前key比较小，需要放置右边
            else {
                if (nonNull(this.right)) {
                    this.right.insert(node, tree);
                } else {
                    this.right = node;
                    node.parent = this;
                    node.balanceTree(false, tree);
                }
            }
            return null;
        }

        /**
         * 变色
         */
        public void toggleColor() {
            this.color = !this.color;
        }

        /**
         * 平衡tree<br />
         * 1、双红，且叔叔节点存在; 将父节点和叔叔节点变成黑色，爷爷节点变成红色; 后续处理 <br />
         * 2、左左红，且叔叔节点不存在 <br />
         * 3、左右红，且叔叔节点不存在 <br />
         * 4、右右红，且叔叔节点不存在 <br />
         * 5、右左红，且叔叔节点不存在 <br />
         *
         * @param left 当前节点是不是左子节点
         * @param tree 当前红黑树
         */
        public void balanceTree(boolean left, BRTree<K, V> tree) {
            // 双红
            boolean doubleRed = RED.equals(this.parent.getColor());
            // 叔叔节点是否存在
            BRTreeNode<K, V> grandParentNode = this.parent.parent;
            BRTreeNode<K, V> uncleNode = null;
            boolean existsUncle = nonNull(grandParentNode) && nonNull(uncleNode = left ? grandParentNode.getRight() : grandParentNode.getLeft());

            // 双红，且叔叔节点存在
            if (doubleRed && existsUncle) {
                this.toggleTreeColor(uncleNode, tree);
            } else if (doubleRed) {
                boolean parentLeft = grandParentNode.getLeft() == this.parent;
                // 左左红
                if (parentLeft && left) {
                    leftLeftRed(tree);
                }
                // 左右红
                else if (parentLeft) {
                    // 先左旋，变成左左红的情况
                    tree.leftRotate(this.parent);
                    this.left.leftLeftRed(tree);
                }
                // 右右红
                else if (!left) {
                    rightRightRed(tree);
                }
                // 右左红
                else {
                    // 先右旋，变成右右红的情况
                    tree.rightRotate(this.parent);
                    this.right.rightRightRed(tree);
                }
            }
        }

        /**
         * 变色 <br />
         * 1、将父节点和叔叔节点变成黑色 <br />
         * 2、爷爷节点变成红色 <br />
         * 3、递归后续处理 <br />
         *
         * @param uncleNode 叔叔节点
         * @param tree 当前的红黑树
         */
        private void toggleTreeColor(BRTreeNode<K, V> uncleNode, BRTree<K, V> tree) {
            this.parent.toggleColor();
            uncleNode.toggleColor();
            BRTreeNode<K, V> grandParentNode = this.parent.parent;
            grandParentNode.toggleColor();
            // 查看爷爷节点是不是根节点
            if (Objects.isNull(grandParentNode.parent)) {
                // 需要重新变为黑色
                grandParentNode.toggleColor();
            } else {
                // 递归处理后续
                grandParentNode.balanceTree(grandParentNode.parent.getLeft() == grandParentNode, tree);
            }
        }

        /**
         * 左左红 <br />
         * 1、将父节点变成黑色，爷爷节点变成红色 <br />
         * 2、将爷爷节点进行右旋 <br />
         * 3、递归后续处理 <br />
         */
        private void leftLeftRed(BRTree<K, V> tree) {
            // 将父节点变成黑色，爷爷节点变成红色
            this.parent.toggleColor();
            BRTreeNode<K, V> grandParentNode = this.parent.parent;
            grandParentNode.toggleColor();
            // 将爷爷节点进行右旋
            tree.rightRotate(grandParentNode);
            // 递归后续处理
            if (nonNull(grandParentNode.parent)) {
                grandParentNode.balanceTree(grandParentNode.parent.getLeft() == grandParentNode, tree);
            }
        }

        /**
         * 右右红 <br />
         * 1、变色，父节点变成黑色，爷爷节点变成红色 <br />
         * 2、将爷爷节点进行左旋 <br />
         * 3、递归后续处理 <br />
         *
         * @param tree 当前红黑树
         */
        private void rightRightRed(BRTree<K, V> tree) {
            // 将父节点变成黑色，爷爷节点变成红色
            this.parent.toggleColor();
            BRTreeNode<K, V> grandParentNode = this.parent.parent;
            grandParentNode.toggleColor();
            // 将爷爷节点进行左旋
            tree.leftRotate(grandParentNode);
            // 递归后续处理
            if (nonNull(grandParentNode.parent)) {
                grandParentNode.balanceTree(grandParentNode.parent.getLeft() == grandParentNode, tree);
            }
        }

    }

}